import copy
import math
import time
from collections import Counter

import matplotlib.pyplot as plt
import nltk
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
from nltk import word_tokenize
from torch.autograd import Variable

from langconv import Converter

PAD = 0  # padding占位符的索引
UNK = 1  # 未登录词标识符的索引
BOS = 2  # 词标识符的开始标志
EOS = 3  # 词标识符的结束标志
BATCH_SIZE = 700  # 批次大小
EPOCHS = 200  # 训练轮数
LAYERS = 6  # transformer中encoder、decoder层数
H_NUM = 8  # 多头注意力个数
D_MODEL = 256  # 输入、输出词向量维数
D_FF = 1024  # feed forward全连接层维数
DROPOUT = 0.1  # dropout比例
MAX_LENGTH = 60  # 语句最大长度

TRAIN_FILE = 'nmt/en-cn/train_mini_mini_simple.txt'  # 训练集
DEV_FILE = "nmt/en-cn/test_mini_mini_simple.txt"  # 验证集
SAVE_FILE = 'save/last.pt'  # 模型保存路径
DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

print(f'当前device:{DEVICE}')


def seq_padding(X, padding=PAD):
    """
    按批次（batch）对数据填充、长度对齐
    """
    # 计算该批次各条样本语句长度
    L = [len(x) for x in X]
    # 获取该批次样本中语句长度最大值
    ML = max(L)
    # 遍历该批次样本，如果语句长度小于最大长度，则用padding填充
    return np.array([
        np.concatenate([x, [padding] * (ML - len(x))]) if len(x) < ML else x for x in X
    ])


def cht_to_chs(sent):
    '''繁体字转为简体
    Args:
        sent: 简体字

    Returns: 繁体字

    '''
    sent = Converter("zh-hans").convert(sent)
    sent.encode("utf-8")
    return sent


class PrepareData:
    # 类变量，用于保存唯一实例
    _instance: "PrepareData" = None

    def __new__(cls, *args):
        if cls._instance is None:
            cls._instance = super(PrepareData, cls).__new__(cls)
        return cls._instance

    def __init__(self, train_file, dev_file):
        # 读取数据、分词
        self.train_en, self.train_cn = self.load_data(train_file)
        self.dev_en, self.dev_cn = self.load_data(dev_file)
        # 构建词表
        self.en_word_dict, self.en_total_words, self.en_index_dict = \
            self.build_dict(self.train_en)
        self.cn_word_dict, self.cn_total_words, self.cn_index_dict = \
            self.build_dict(self.train_cn)
        # 单词映射为索引
        self.train_en, self.train_cn = self.word2id(self.train_en, self.train_cn, self.en_word_dict, self.cn_word_dict)
        self.dev_en, self.dev_cn = self.word2id(self.dev_en, self.dev_cn, self.en_word_dict, self.cn_word_dict)
        # 划分批次、填充、掩码
        self.train_data = self.split_batch(self.train_en, self.train_cn, BATCH_SIZE)
        self.dev_data = self.split_batch(self.dev_en, self.dev_cn, BATCH_SIZE)

    def load_data(self, path):
        """
        读取英文、中文数据
        对每条样本分词并构建包含起始符和终止符的单词列表
        形式如：en = [['BOS', 'i', 'love', 'you', 'EOS'], ['BOS', 'me', 'too', 'EOS'], ...]
               cn = [['BOS', '我', '爱', '你', 'EOS'], ['BOS', '我', '也', '是', 'EOS'], ...]
        """
        en = []
        cn = []
        with open(path, mode="r", encoding="utf-8") as f:
            for line in f.readlines():
                sent_en, sent_cn = line.strip().split("\t")
                sent_en = sent_en.lower()
                # 繁体字转简体字
                # sent_cn = cht_to_chs(sent_cn)
                sent_en = ["BOS"] + word_tokenize(sent_en) + ["EOS"]
                # 中文按字符切分
                sent_cn = ["BOS"] + [char for char in sent_cn] + ["EOS"]
                en.append(sent_en)
                cn.append(sent_cn)
        return en, cn

    def build_dict(self, sentences, max_words=5e4):
        '''
        构建映射字典
        :param sentences: 输入[['BOS', 'i', 'love', 'you', 'EOS'], ['BOS', 'me', 'too', 'EOS'], ...]
        :param max_words: 最大映射字典数
        :return: word_dict: {'PAD': 0, 'UNK': 1, 'BOS': 2,'EOS': 3,'.': 4,'i': 5,'the': 6,'to': 7,'you': 8}
                 total_words: 字典大小
                 index_dict: {0: 'PAD', 1: 'UNK', 2: 'BOS', 3: 'EOS', 4: '.', 5: 'i', 6: 'the', 7: 'to', 8: 'you'}
        '''
        # 统计数据集中单词词频
        word_count = Counter([word for sent in sentences for word in sent])
        # 按词频保留前max_words个单词构建词典
        ls = word_count.most_common(int(max_words))
        # 添加UNK和PAD两个单词
        total_words = len(ls) + 2
        word_dict = {
            'UNK': UNK,  # 不在句子中
            'PAD': PAD,  # 不在句子中
            'BOS': BOS,  # 在句子中
            'EOS': EOS,  # 在句子中
        }
        # 将'BOS'和'EOS'从ls中移除
        ls = [item for item in ls if item[0] not in ('BOS', 'EOS')]
        for index, w in enumerate(ls):
            word_dict.update({w[0]: index + 4})

        # 根据word_dict的value进行升序排序
        word_dict = dict(sorted(word_dict.items(), key=lambda item: item[1]))
        # 构建id2word映射
        index_dict = {v: k for k, v in word_dict.items()}
        return word_dict, total_words, index_dict

    def word2id(self, en, cn, en_dict, cn_dict, sort=True):
        """
        :argument en:[['BOS', 'anyone', 'can', 'do', 'that', '.', 'EOS']]
        :argument cn:[['BOS', '任', '何', '人', '都', '可', '以', '做', '到', '。', 'EOS']]
        :argument en_dict:{'PAD': 0, 'UNK': 1, 'BOS': 2,'EOS': 3,'.': 4,'i': 5,'the': 6,'to': 7,'you': 8}
        :argument cn_dict:{'PAD': 0, 'UNK': 1, 'BOS': 2,'EOS': 3,'。': 4,'我': 5,'的': 6,'了': 7,'你': 8}
        :return
                out_en_ids: [ [ 2, 475, 4, 3 ], [ 2, 1318, 126, 3 ]]
                out_cn_ids: [ [ 2, 8, 90, 427, 6, 4, 3 ], [ 2, 127, 871, 217, 3 ]]

        将英文、中文单词列表转为单词索引列表
        `sort=True`表示以英文语句长度排序，以便按批次填充时，同批次语句填充尽量少
        """
        # 单词映射为索引
        out_en_ids = [[en_dict.get(word, UNK) for word in sent] for sent in en]
        out_cn_ids = [[cn_dict.get(word, UNK) for word in sent] for sent in cn]

        # 按照语句长度排序
        def len_argsort(seq):
            """
            传入一系列语句数据(分好词的列表形式)，
            按照语句长度排序后，返回排序后原来各语句在数据中的索引下标
            """
            return sorted(range(len(seq)), key=lambda x: len(seq[x]))

        # 按相同顺序对中文、英文样本排序
        if sort:
            # 以英文语句长度排序(升序)
            sorted_index = len_argsort(out_en_ids)
            out_en_ids = [out_en_ids[idx] for idx in sorted_index]
            out_cn_ids = [out_cn_ids[idx] for idx in sorted_index]
        return out_en_ids, out_cn_ids

    def split_batch(self, en, cn, batch_size, shuffle=True):
        '''

        Args:
            en: [ [ 2, 475, 4, 3 ], [ 2, 1318, 126, 3 ] 每个英文语句的id
            cn:[ [ 2, 8, 90, 427, 6, 4, 3 ], [ 2, 127, 871, 217, 3 ]] 每个中文语句的id
            batch_size:
            shuffle: 是否随机打乱

        Returns:

        '''
        # 每隔batch_size取一个索引作为后续batch的起始索引
        idx_list = np.arange(0, len(en), batch_size)
        # 起始索引随机打乱
        if shuffle:
            np.random.shuffle(idx_list)
        # 存放所有批次的语句索引
        batch_indexs = []
        for idx in idx_list:
            """
            形如[array([4, 5, 6, 7]), 
                 array([0, 1, 2, 3]), 
                 array([8, 9, 10, 11]),
                 ...]
            """
            # 起始索引最大的批次可能发生越界，要限定其索引
            batch_indexs.append(np.arange(idx, min(idx + batch_size, len(en))))
        # 构建批次列表
        batches = []
        for batch_index in batch_indexs:
            # 按当前批次的样本索引采样
            batch_en = [en[index] for index in batch_index]
            batch_cn = [cn[index] for index in batch_index]
            # 对当前批次中所有语句填充、对齐长度
            # 维度为：batch_size * 当前批次中语句的最大长度
            batch_cn = seq_padding(batch_cn)
            batch_en = seq_padding(batch_en)
            # 将当前批次添加到批次列表
            # Batch类用于实现注意力掩码
            batches.append(Batch(batch_en, batch_cn))
        return batches


# 示例
# data_mini = PrepareData("./nmt/en-cn/train_mini.txt", "./nmt/en-cn/dev_mini.txt")
# print("*** batch 0, en ***")
# print(data_mini.train_data[0][0])
# print("*** batch 0, chs ***")
# print(data_mini.train_data[0][1])
# print("*** en word to index ***")
# print(data_mini.en_word_dict)
# print("*** chs word to index ***")
# print(data_mini.cn_word_dict)
class Embeddings(nn.Module):
    def __init__(self, d_model, vocab):
        super(Embeddings, self).__init__()
        # Embedding层
        self.lut = nn.Embedding(vocab, d_model)
        # Embedding维数
        self.d_model = d_model

    def forward(self, x):
        # 返回x的词向量（需要乘以math.sqrt(d_model)）
        return self.lut(x) * math.sqrt(self.d_model)


class PositionalEncoding(nn.Module):
    def __init__(self, d_model, dropout, max_len=5000):
        super(PositionalEncoding, self).__init__()
        self.dropout = nn.Dropout(p=dropout)
        # 位置编码矩阵，维度[max_len, embedding_dim]
        pe = torch.zeros(max_len, d_model, device=DEVICE)
        # 单词位置
        position = torch.arange(0.0, max_len, device=DEVICE)
        position.unsqueeze_(1)
        # 使用exp和log实现幂运算
        div_term = torch.exp(torch.arange(0.0, d_model, 2, device=DEVICE) * (- math.log(1e4) / d_model))
        div_term.unsqueeze_(0)
        # 计算单词位置沿词向量维度的纹理值
        pe[:, 0:: 2] = torch.sin(torch.mm(position, div_term))
        pe[:, 1:: 2] = torch.cos(torch.mm(position, div_term))
        # 增加批次维度，[1, max_len, embedding_dim]
        pe.unsqueeze_(0)
        # 将位置编码矩阵注册为buffer(不参加训练)
        self.register_buffer('pe', pe)

    def forward(self, x):
        # 将一个批次中语句所有词向量与位置编码相加
        # 注意，位置编码不参与训练，因此设置requires_grad=False
        x += Variable(self.pe[:, : x.size(1), :], requires_grad=False)
        return self.dropout(x)


# emb_dim = 64
# max_seq_len = 100
# seq_len = 20

# pe = PositionalEncoding(emb_dim, 0, max_seq_len)
# positional_encoding = pe(torch.zeros(1, seq_len, emb_dim, device=DEVICE))
# plt.figure()
# sns.heatmap(positional_encoding.squeeze().to("cpu"))
# plt.xlabel("i")
# plt.ylabel("pos")
# plt.show()
#
# plt.figure()
# y = positional_encoding.to("cpu").numpy()
# plt.plot(np.arange(seq_len), y[0, :, 0 : 64 : 8], ".")
# plt.legend(["dim %d" % p for p in [0, 7, 15, 31, 63]])
# plt.show()

def clones(module, N):
    """
    克隆基本单元，克隆的单元之间参数不共享
    """
    return nn.ModuleList([
        copy.deepcopy(module) for _ in range(N)
    ])


def attention(query, key, value, mask=None, dropout=None):
    '''

    :param query: [batch_size,head_num,seq_len,每个头的维度]
    :param key:  [batch_size,head_num,seq_len,每个头的维度]
    :param value:  [batch_size,head_num,seq_len,每个头的维度]
    :param mask: [batch_size,1,1,seq_len],用于遮蔽某些位置，通常用于填充的部分,为0的部分，填充-1e9
    :param dropout: 可选模块，Dropout层，用于在注意力计算后进行随机失活
    :return:
        attn_output:[batch_size,head_num,seq_len,seq_len] 加权后的输出
        attn_weights:[batch_size,head_num,seq_len,seq_len] 注意力权重矩阵
    '''
    # q、k、v向量长度为d_k
    d_k = query.size(-1)
    # 矩阵乘法实现q、k点积注意力，sqrt(d_k)归一化
    # [batch_size,head_num,seq_len,seq_len]
    scores = (query @ key.transpose(-2, -1)) / math.sqrt(d_k)
    # 注意力掩码机制
    if mask is not None:
        '''
        用-1e9代替0 -1e9是一个很大的负数 经过softmax之后接近0  
        其一：去除掉各种padding在训练过程中的影响  
        其二，将输入进行遮盖，避免decoder看到后面要预测的东西。（只用在decoder中）
        '''
        scores = scores.masked_fill(mask == 0, -1e9)
    # 注意力矩阵softmax归一化
    p_attn = F.softmax(scores, dim=-1)
    # dropout
    if dropout is not None:
        p_attn = dropout(p_attn)
    # 注意力对v加权
    return (p_attn @ value), p_attn


class MultiHeadedAttention(nn.Module):
    def __init__(self, num_heads, d_model, dropout=0.1):
        super(MultiHeadedAttention, self).__init__()
        """
        `h`：注意力头的数量
        `d_model`：词向量维数
        """
        # 确保整除
        assert d_model % num_heads == 0
        # q、k、v向量每个头的维数
        self.d_k = d_model // num_heads
        # 头的数量
        self.num_heads = num_heads
        # WQ、WK、WV矩阵及多头注意力拼接变换矩阵WO
        self.linears = clones(nn.Linear(d_model, d_model), 4)
        self.attn = None
        self.dropout = nn.Dropout(p=dropout)

    def forward(self, query, key, value, mask=None):
        '''

        :param query: [batch_size, seq_len, d_model]
        :param key: [batch_size, seq_len, d_model]
        :param value: [batch_size, seq_len, d_model]
        :param mask: [batch_size, 1, seq_len] 用于屏蔽某些位置的注意力
        :return:    [batch_size, seq_len, d_model]
        '''
        if mask is not None:
            # 为head增加一个维度
            mask = mask.unsqueeze(1)
        # 批次大小
        nbatches = query.size(0)
        # WQ、WK、WV分别对词向量线性变换，并将结果拆成h块(拆分为多头)
        # [batch_size,head_nums,seq_len,d_k]
        query, key, value = [
            l(x).view(nbatches, -1, self.num_heads, self.d_k).transpose(1, 2)
            for l, x in zip(self.linears, (query, key, value))
        ]
        # 注意力加权
        x, self.attn = attention(query, key, value, mask=mask, dropout=self.dropout)
        # 多头注意力加权拼接
        # [batch_size,head_nums,seq_len,d_k] -> [batch_size,seq_len,d_model]
        x = x.transpose(1, 2).contiguous().view(nbatches, -1, self.num_heads * self.d_k)
        # 对多头注意力加权拼接结果线性变换
        out = self.linears[-1](x)
        return out


# 前馈神经网络
class PositionwiseFeedForward(nn.Module):
    def __init__(self, d_model, d_ff, dropout=0.1):
        super(PositionwiseFeedForward, self).__init__()
        self.w_1 = nn.Linear(d_model, d_ff)
        self.w_2 = nn.Linear(d_ff, d_model)
        self.dropout = nn.Dropout(dropout)

    def forward(self, x):
        x = self.w_1(x)
        x = F.relu(x)
        x = self.dropout(x)
        x = self.w_2(x)
        return x


# 单个编码层
class EncoderLayer(nn.Module):
    def __init__(self, d_model=512, d_ff=2048, num_heads=8, dropout=0.1):
        '''

        :param d_model:  模型的维度
        :param d_ff: 前馈网络中间层的维度
        :param num_heads: 前馈神经网络
        :param dropout: 置零比例
        '''
        super(EncoderLayer, self).__init__()
        # 多头自注意力层
        self.self_attn = MultiHeadedAttention(num_heads, d_model)
        # 前馈神经网络层
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff, dropout)
        # 第一个层归一化
        self.norm1 = nn.LayerNorm(d_model)
        # 第二个层归一化
        self.norm2 = nn.LayerNorm(d_model)
        # Dropout层
        self.dropout = nn.Dropout(dropout)  # Dropout层

    def forward(self, x, mask):
        '''编码器层的前向传播
        :param x: [batch_size, seq_len, d_model]
        :param mask:[1,1,seq_len]
        :return: 处理后的张量 [batch_size, seq_len, d_model]
        '''
        # 多头自注意力子层（带有残差连接和层归一化）
        attn_output = self.self_attn(x, x, x, mask)  # 自注意力：查询、键、值都是x
        x = self.norm1(x + self.dropout(attn_output))  # 残差连接和层归一化

        # 前馈网络子层（带有残差连接和层归一化）
        ff_output = self.feed_forward(x)
        x = self.norm2(x + self.dropout(ff_output))  # 残差连接和层归一化
        return x


# 编码器
class Encoder(nn.Module):
    def __init__(self, num_layers=6, d_model=512, d_ff=2048, num_heads=8, dropout=0.1):
        super(Encoder, self).__init__()
        # 创建多个编码器层
        self.layers = nn.ModuleList([EncoderLayer(d_model, d_ff, num_heads, dropout) for _ in range(num_layers)])

    def forward(self, x, mask):
        '''

        :param x: [batch_size, seq_len, d_model]
        :param mask: [1, 1, seq_len]
        :return:
        '''
        # 循环编码器基本单元N次
        for layer in self.layers:
            x = layer(x, mask)
        return x


# 解码器
class Decoder(nn.Module):
    def __init__(self, num_layers, d_model, d_ff, num_heads, dropout):
        super(Decoder, self).__init__()
        # 创建多个解码器层
        self.layers = nn.ModuleList([DecoderLayer(d_model, d_ff, num_heads, dropout) for _ in range(num_layers)])

    def forward(self, x, memory, src_mask, trg_mask):
        """
        循环解码器基本单元N次
        """
        for layer in self.layers:
            x = layer(x, memory, src_mask, trg_mask)
        return x


# 单个解码器层
class DecoderLayer(nn.Module):
    def __init__(self, d_model, d_ff, num_heads, dropout):
        super(DecoderLayer, self).__init__()
        # 自注意力层
        self.self_attn = MultiHeadedAttention(num_heads, d_model)
        # 交叉注意力层
        self.cross_attn = MultiHeadedAttention(num_heads, d_model)
        # 前馈神经网络层
        self.feed_forward = PositionwiseFeedForward(d_model, d_ff)
        self.norm1 = nn.LayerNorm(d_model)  # 第一个层归一化
        self.norm2 = nn.LayerNorm(d_model)  # 第二个层归一化
        self.norm3 = nn.LayerNorm(d_model)  # 第三个层归一化
        self.dropout = nn.Dropout(dropout)  # Dropout层

    def forward(self, x, memory, src_mask, trg_mask):
        '''
        解码器层的前向传播
        Args:
            x: 目标语言输入张量 [batch_size, seq_len, d_model]
            memory: 编码器输出 [batch_size, src_seq_len, d_model]
            src_mask: 源序列掩码
            trg_mask: 目标序列掩码（防止看到未来位置）

        Returns: 处理后的张量 [batch_size, seq_len, d_model]

        '''
        # 掩码多头自注意力子层（带有残差连接和层归一化）
        attn_output = self.self_attn(x, x, x, trg_mask)  # 自注意力：查询、键、值都是x
        x = self.norm1(x + self.dropout(attn_output))  # 残差连接和层归一化
        # 编码器-解码器多头注意力子层（带有残差连接和层归一化）
        # 查询来自解码器的前一个子层，键和值来自编码器的输出
        attn_output = self.cross_attn(x, memory, memory, src_mask)
        x = self.norm2(x + self.dropout(attn_output))  # 残差连接和层归一化
        # 前馈网络子层（带有残差连接和层归一化）
        ff_output = self.feed_forward(x)
        x = self.norm3(x + self.dropout(ff_output))  # 残差连接和层归一化
        return x


class Generator(nn.Module):
    """
    该模块的作用是把解码器输出的隐向量从d_model维度映射到词表长度，得到logits
    logits对应着该token取不同字的概率，接下来模型会依据这些概率，按照一定的采样规则来采样下一个token。
    模型的效果好坏就是看看模型是否可以把下一个token分类到真值对应的token。
    因此，每次预测新token都是一个分类任务，Generator就是一个分类头。训练会依据分类结果来计算损失。
    """

    def __init__(self, d_model, vocab):
        super(Generator, self).__init__()
        # decode后的结果，先进入一个全连接层变为词典大小的向量
        self.proj = nn.Linear(d_model, vocab)

    def forward(self, x):
        # 然后再进行log_softmax操作(在softmax结果上再做多一次log运算)
        return F.log_softmax(self.proj(x), dim=-1)


def subsequent_mask(size):
    '''
    下三角矩阵，生成解码器的掩码矩阵，用于屏蔽未来位置的信息
    Args:
        size:  序列长度

    Returns:  掩码矩阵，形状为 [1,sz, sz]

    '''
    # 设定subsequent_mask矩阵的shape
    attn_shape = (1, size, size)
    # 生成一个右上角(不含主对角线)为全1，左下角(含主对角线)为全0的subsequent_mask矩阵
    subsequent_mask = np.triu(np.ones(attn_shape), k=1).astype('uint8')
    # 返回一个右上角(不含主对角线)为全False，左下角(含主对角线)为全True的subsequent_mask矩阵
    return torch.from_numpy(subsequent_mask) == 0


class Batch:
    """
    批次类
        1. 输入序列（源）
        2. 输出序列（目标）
        3. 构造掩码
    """

    def __init__(self, src, trg=None, pad=PAD):
        '''

        Args:
            src: [batch_size,max_len_en] 经过填充的英语id序列
            trg: [batch_size,max_len_cn] 经过填充的汉语id序列
            pad: 填充id
        '''

        # 将输入、输出单词id表示的数据规范成整数类型
        src = torch.from_numpy(src).to(DEVICE).long()
        trg = torch.from_numpy(trg).to(DEVICE).long()

        src_0 = src[0].data.numpy().tolist()
        trg_0 = trg[0].data.numpy().tolist()

        en = [PrepareData._instance.en_index_dict.get(item, None) for item in src_0]
        cn = [PrepareData._instance.cn_index_dict.get(item, None) for item in trg_0]

        self.src = src
        # 创建源序列的掩码。填充符号(PAD)的地方会被忽略，因此生成的掩码是一个bool类型的矩阵，填充符号的位置为False
        # 并在seq length前面增加一维，形成维度为[batch_size, 1,seq length] 的矩阵,False代表需要遮掩的部分
        self.src_mask = (src != pad).unsqueeze(-2)
        # 如果输出目标不为空，则需要对解码器使用的目标语句进行掩码
        if trg is not None:
            # 解码器使用的目标输入部分
            self.trg_x = trg[:, :-1]  # 去除最后一列
            # 解码器训练时应预测输出的目标结果
            self.trg_y = trg[:, 1:]  # 去除第一列的BOS
            # 将目标输入部分进行注意力掩码
            self.trg_mask = self.make_std_mask(self.trg_x, pad)
            # 将应输出的目标结果中实际的词数进行统计
            self.ntokens = (self.trg_y != pad).data.sum()

    # 掩码操作
    @staticmethod
    def make_std_mask(trg_x, pad):
        # 创建目标序列的掩码。填充符号(PAD)的地方会被忽略，因此生成的掩码是一个bool类型的矩阵，填充符号的位置为False
        # 并在seq length前面增加一维，形成维度为[batch_size, 1,seq length] 的矩阵,False代表需要遮掩的部分
        tgt_mask = (trg_x != pad).unsqueeze(-2)
        # 合并三角矩阵和pad矩阵
        tgt_mask = tgt_mask & Variable(subsequent_mask(trg_x.size(-1)).type_as(tgt_mask.data))
        return tgt_mask


class Transformer(nn.Module):
    def __init__(self, src_vocab, trg_vocab, num_layers=6, d_model=512, d_ff=2048, num_heads=8, dropout=0.1):
        '''

        Args:
            src_vocab: 源语言词汇表大小
            trg_vocab: 目标语言词汇表大小
            num_layers: 编码器和解码器的层数，默认为6
            d_model:  模型的维度，默认为512
            d_ff: 前馈网络中间层的维度，默认为2048
            num_heads: 注意力头的数量，默认为8
            dropout: dropout比率，默认为0.1
        '''
        super(Transformer, self).__init__()
        # 编码器
        self.encoder = Encoder(num_layers, d_model, d_ff, num_heads, dropout)
        # 解码器
        self.decoder = Decoder(num_layers, d_model, d_ff, num_heads, dropout)

        # 源语言词嵌入层 + 位置编码层
        self.src_embed = nn.Sequential(Embeddings(d_model, src_vocab), PositionalEncoding(d_model, dropout))
        # 目标语言词嵌入层 + 位置编码层
        self.trg_embed = nn.Sequential(Embeddings(d_model, trg_vocab), PositionalEncoding(d_model, dropout))

        # 最终的线性层，将解码器输出映射到词汇表大小
        self.generator = Generator(d_model, trg_vocab)

    def encode(self, src, src_mask):
        '''

        Args:
            src: [batch_size,src_len]
            src_mask: [batch_size,1,src_len]

        Returns: [batch_size,seq_len,d_model]

        '''
        x = self.src_embed(src)
        x = self.encoder(x, src_mask)
        return x

    def decode(self, memory, src_mask, trg_x, trg_mask):
        '''

        Args:
            memory: [batch_size,src_len,d_model] 编码器输出
            src_mask: [batch_size,1,src_len] 源序列掩码
            trg_x:    [batch_size,src_len] 目标序列输入
            trg_mask: [batch_size,1,src_len] 目标序列掩码

        Returns: [batch_size,src_len,d_model]

        '''
        x = self.trg_embed(trg_x)
        x = self.decoder(x, memory, src_mask, trg_mask)
        return x

    def forward(self, src, trg_x, src_mask, trg_mask):
        '''

        Args:
            src: [batch_size,en_seq_len] 源序列输入
            trg_x: [batch_size,cn_seq_len] 目标序列输入
            src_mask: [batch_size,1,en_seq_len] 源序列掩码
            trg_mask: [batch_size,cn_seq_len,cn_seq_len] 目标序列掩码

        Returns: 返回张量 [batch_size,cn_seq_len,d_model]

        '''
        # encoder的结果作为decoder的memory参数传入，进行decode
        x = self.encode(src, src_mask)
        x = self.decode(x, src_mask, trg_x, trg_mask)
        return x


def make_model(src_vocab, trg_vocab, num_layers=6, d_model=512, d_ff=2048, num_heads=8, dropout=0.1):
    # 实例化Transformer模型对象
    model = Transformer(src_vocab, trg_vocab, num_layers, d_model, d_ff, num_heads, dropout).to(DEVICE)

    # This was important from their code.
    # Initialize parameters with Glorot / fan_avg.
    for p in model.parameters():
        if p.dim() > 1:
            # 这里初始化采用的是nn.init.xavier_uniform
            nn.init.xavier_uniform_(p)
    return model


class LabelSmoothing(nn.Module):
    """
    LabelSmoothing类除了负责平滑标签外，还负责计算损失
    size: 目标语言词典大小。
    padding_idx: <pad>在词典中对应的序号
    smoothing: 平滑因子，0表示不做平滑处理
    标签平滑（Label Smoothing）是一种正则化技术，用于深度学习中的分类任务，
    尤其是在Transformer模型中。它的目的是减少模型对于训练数据中硬标签（hard labels，即标准的one-hot编码）的过度自信，
    从而提高模型的泛化能力并减少过拟合。
    例子:
        硬标签[1,0,0],平滑系数ϵ为0.1,
        那么正确标签概率为 1-0.1=0.9
        其他类别标签概率为 0.1/3=0.0333
        因此，平滑后的软标签为[0.9,0.0333,0.0333]
    """

    def __init__(self, size, padding_idx, smoothing=0.0):
        '''

        Args:
            size: 标语言词典大小。
            padding_idx: <pad>在词典中对应的序号
            smoothing: 平滑因子，0表示不做平滑处理
        '''
        super(LabelSmoothing, self).__init__()
        '''
        在 Transformer 翻译模型中使用 KL散度 (Kullback-Leibler Divergence)，主要是因为它能够衡量模型生成的概率分布与目标分布之间的差异
        '''
        self.criterion = nn.KLDivLoss(reduction='sum')
        self.padding_idx = padding_idx
        self.confidence = 1.0 - smoothing
        self.smoothing = smoothing
        self.size = size
        self.true_dist = None  # 平滑后的标签

    def forward(self, x, target):
        '''

        Args:
            x: generator输出的概率分布。形状为(在一个批次中的单词个数, voc_size)
            target: 目标真值标签，内容是token index。形状为(在一个批次中的单词个数)

        Returns:

        '''
        assert x.size(1) == self.size
        # 创建一个与x有相同形状的张量,最后一个维度代表所在序列的目标词汇概率
        true_dist = x.data.clone()
        # 将true_dist全部填充为 self.smoothing / (self.size - 2)
        """
            假设 smoothing=0.2，词表大小为6，batch size为2
            则true_dist全部填充为 0.2 / (6-2)= 0.05，此时true_dist为：
            [[0.05, 0.05, 0.05, 0.05, 0.05, 0.05],
             [0.05, 0.05, 0.05, 0.05, 0.05, 0.05]]
        """
        true_dist.fill_(self.smoothing / (self.size - 2))
        '''
             target.data.unsqueeze(1)会给target.data增加一维，假设target.data是[2,3]，则target.data.unsqueeze(1)的结果是[[2],[3]]
             将true_dist的第一个1维度上与target.data.unsqueeze(1)对应的值变为self.confidence。
             假设此例中target.data.unsqueeze(1) 为[[2], [3]]，即2个数据的标签分别为2，3，就是把true_dist上设置为self.confidence，
             则true_dist执行过scatter后变为:
             [[0.05, 0.05, 0.8, 0.05, 0.05, 0.05],
             [0.05, 0.05, 0.05, 0.8, 0.05, 0.05]]
        '''
        true_dist.scatter_(1, target.data.unsqueeze(1), self.confidence)
        # 将<pad>所在的index填充为0，pad填充字符不需要计算损失
        true_dist[:, self.padding_idx] = 0
        # 找出target中为<pad>的标签。例如target为['i', 'love', 'you', '<pad>', '<pad>']，mask则为[[3], [4]]，表示第3个和第4个为空格。
        mask = torch.nonzero(target.data == self.padding_idx)
        if mask.dim() > 0:
            # 将"<pad>"所在的label设置为0
            true_dist.index_fill_(0, mask.squeeze(), 0.0)
        # 保存平滑标签后的label
        self.true_dist = true_dist
        """
            使用平滑后的标签计算损失
            由于对`<pad>`部分进行了mask，所以这部分不会参与损失计算
        """
        return self.criterion(x, Variable(true_dist, requires_grad=False))


class SimpleLossCompute:
    """
    简单的计算损失和进行参数反向传播更新训练的函数
    """

    def __init__(self, generator: Generator, criterion: LabelSmoothing, opt=None):
        self.generator = generator
        self.criterion = criterion
        self.opt = opt

    def __call__(self, x, y, norm):
        '''
        假设有两个batch，第一个batch有6个字，则loss是这6个预测结果计算损失的和。
        第二个batch有60个字，则loss是这60个预测结果计算损失的和。
        显然第二个损失大，这不符合逻辑。所以我们用除以有效token数目来进行平均。
        Args:
            x:  解码器的输出 [batch_size,seq_len,d_model]
            y:  batch.tgt_y，要被预测的所有token，例如src为`<bos>我吃了一个苹果<eos>`，
                则tgt_y是"I ate an apple<eos>" [batch_size,seq_len]
            norm: batch.ntokens, tgt_y中的有效token数

        Returns:

        '''
        # [batch_size,seq_len,d_model] --》 [batch_size,seq_len,tgr_vocab]
        x = self.generator(x)  # 生成预测输出
        # 首先使用KLDivLoss进行了损失计算，随后又除以batch.ntokens对损失进行正则化
        # [batch_size*seq_len,tgr_vocab]，转为计算总的单词在目标词汇上的输出概率
        loss = self.criterion(x.contiguous().view(-1, x.size(-1)), y.contiguous().view(-1)) / norm
        loss.backward()
        if self.opt is not None:
            self.opt.step()
            self.opt.optimizer.zero_grad()
        return loss.data.item() * norm.float()


class NoamOpt:
    "Optim wrapper that implements rate."

    def __init__(self, model_size, factor, warmup, optimizer):
        self.optimizer = optimizer
        self._step = 0
        self.warmup = warmup
        self.factor = factor
        self.model_size = model_size
        self._rate = 0

    def step(self):
        "Update parameters and rate"
        self._step += 1
        rate = self.rate()
        for p in self.optimizer.param_groups:
            p['lr'] = rate
        self._rate = rate
        self.optimizer.step()

    def rate(self, step=None):
        "Implement `lrate` above"
        if step is None:
            step = self._step
        return self.factor * (self.model_size ** (-0.5) * min(step ** (-0.5), step * self.warmup ** (-1.5)))


def run_epoch(data, model, loss_compute, epoch):
    start = time.time()
    total_tokens = 0.
    total_loss = 0.
    tokens = 0.

    for i, batch in enumerate(data):
        out = model(batch.src, batch.trg_x, batch.src_mask, batch.trg_mask)
        loss = loss_compute(out, batch.trg_y, batch.ntokens)

        total_loss += loss
        total_tokens += batch.ntokens
        tokens += batch.ntokens

        if i % 50 == 1:
            elapsed = time.time() - start
            print("Epoch %d Batch: %d Loss: %f Tokens per Sec: %fs" % (
                epoch, i - 1, loss / batch.ntokens, (tokens.float() / elapsed / 1000.)))
            start = time.time()
            tokens = 0

    return total_loss / total_tokens


def train(data, model, criterion, optimizer):
    """
    训练并保存模型
    """
    # 初始化模型在dev集上的最优Loss为一个较大值
    best_dev_loss = 1e5

    for epoch in range(EPOCHS):
        # 模型训练
        model.train()
        run_epoch(data.train_data, model, SimpleLossCompute(model.generator, criterion, optimizer), epoch)
        model.eval()

        # 在dev集上进行loss评估
        print('>>>>> Evaluate')
        dev_loss = run_epoch(data.dev_data, model, SimpleLossCompute(model.generator, criterion, None), epoch)
        print('<<<<< Evaluate loss: %f' % dev_loss)

        # 如果当前epoch的模型在dev集上的loss优于之前记录的最优loss则保存当前模型，并更新最优loss值
        if dev_loss < best_dev_loss:
            torch.save(model.state_dict(), SAVE_FILE)
            best_dev_loss = dev_loss
            print('****** Save best model done... ******\r\n')


def greedy_decode(model: Transformer, src, src_mask, max_len, start_symbol):
    """
    传入一个训练好的模型，对指定数据进行预测
    """
    # 先用encoder进行encode
    memory = model.encode(src, src_mask)
    # 初始化预测内容为1×1的tensor，填入开始符('BOS')的id，并将type设置为输入数据类型(LongTensor)
    trg_x = torch.ones(1, 1).fill_(start_symbol).type_as(src.data)
    # 遍历输出的长度下标
    for i in range(max_len - 1):
        # decode得到隐层表示
        trg_mask = subsequent_mask(trg_x.size(1)).type_as(src.data)
        out = model.decode(memory,
                           src_mask,
                           trg_x,
                           trg_mask
                           )
        # 将隐藏表示转为对词典各词的log_softmax概率分布表示，取最后一句话预测结果
        prob = model.generator(out[:, -1])
        # 获取当前位置最大概率的预测词id
        _, next_word = torch.max(prob, dim=1)
        next_word = next_word.data[0]

        # 将当前位置预测的字符id与之前的预测内容拼接起来
        trg_x = torch.cat([trg_x,
                           torch.ones(1, 1).type_as(src.data).fill_(next_word)], dim=1)

        # 检查是否已经翻译到中文的结束词了
        if next_word == EOS:
            return trg_x

    return trg_x


def evaluate(data, model):
    """
    在data上用训练好的模型进行预测，打印模型翻译结果
    """
    # 梯度清零
    with torch.no_grad():
        # 在data的英文数据长度上遍历下标
        for i in range(len(data.dev_en)):
            # 打印待翻译的英文语句
            en_sent = " ".join([data.en_index_dict[w] for w in data.dev_en[i]])
            print("\n" + en_sent)

            # 打印对应的中文语句答案
            cn_sent = " ".join([data.cn_index_dict[w] for w in data.dev_cn[i]])
            print("".join(cn_sent))

            # 将当前以单词id表示的英文语句数据转为tensor，并放如DEVICE中
            src = torch.from_numpy(np.array(data.dev_en[i])).long().to(DEVICE)
            # 增加一维
            src = src.unsqueeze(0)
            # 设置attention mask
            src_mask = (src != 0).unsqueeze(-2)
            # 用训练好的模型进行decode预测
            out = greedy_decode(model, src, src_mask, max_len=MAX_LENGTH, start_symbol=data.cn_word_dict["BOS"])
            # 初始化一个用于存放模型翻译结果语句单词的列表
            translation = []
            # 遍历翻译输出字符的下标（注意：开始符"BOS"的索引0不遍历）
            for j in range(1, out.size(1)):
                # 获取当前下标的输出字符
                sym = data.cn_index_dict[out[0, j].item()]
                # 如果输出字符不为'EOS'终止符，则添加到当前语句的翻译结果列表
                if sym != 'EOS':
                    translation.append(sym)
                # 否则终止遍历
                else:
                    break
            # 打印模型翻译输出的中文语句结果
            print("translation: %s" % " ".join(translation))


if __name__ == '__main__':
    # 下载自然语言处理库
    nltk.download('punkt')

    emb_dim = 64
    max_seq_len = 100
    seq_len = 20

    pe = PositionalEncoding(emb_dim, 0, max_seq_len)
    positional_encoding = pe(torch.zeros(1, seq_len, emb_dim, device=DEVICE))

    # plt.figure()
    # sns.heatmap(positional_encoding.squeeze().to("cpu"))
    # plt.xlabel("i")
    # plt.ylabel("pos")
    # plt.show()
    #
    # plt.figure()
    # y = positional_encoding.to("cpu").numpy()
    # plt.plot(np.arange(seq_len), y[0, :, 0: 64: 8], ".")
    # plt.legend(["dim %d" % p for p in [0, 7, 15, 31, 63]])
    # plt.show()
    #
    # plt.figure(figsize=(5, 5))
    # plt.imshow(subsequent_mask(20)[0])
    # plt.show()

    # Label smoothing的例子
    # crit = LabelSmoothing(5, 0, 0.4)  # 设定一个ϵ=0.4
    # predict = torch.FloatTensor([[0, 0.2, 0.7, 0.1, 0],
    #                              [0, 0.2, 0.7, 0.1, 0],
    #                              [0, 0.2, 0.7, 0.1, 0]])
    # v = crit(Variable(predict.log()),
    #          Variable(torch.LongTensor([2, 1, 0])))
    #
    # # Show the target distributions expected by the system.
    # print(crit.true_dist)
    # plt.imshow(crit.true_dist)

    # Three settings of the lrate hyperparameters.
    opts = [NoamOpt(512, 1, 4000, None),
            NoamOpt(512, 1, 8000, None),
            NoamOpt(256, 1, 4000, None)]
    plt.plot(np.arange(1, 20000), [[opt.rate(i) for opt in opts] for i in range(1, 20000)])
    plt.legend(["512:4000", "512:8000", "256:4000"])

    # 数据预处理
    data = PrepareData(TRAIN_FILE, DEV_FILE)

    src_vocab = len(data.en_word_dict)
    trg_vocab = len(data.cn_word_dict)
    print("src_vocab %d" % src_vocab)
    print("tgt_vocab %d" % trg_vocab)

    # 初始化模型
    model = make_model(
        src_vocab,
        trg_vocab,
        LAYERS,
        D_MODEL,
        D_FF,
        H_NUM,
        DROPOUT
    )

    # 训练
    print(">>>>>>> start train")
    train_start = time.time()
    criterion = LabelSmoothing(trg_vocab, padding_idx=0, smoothing=0.0)
    optimizer = NoamOpt(D_MODEL, 1, 2000, torch.optim.Adam(model.parameters(), lr=0, betas=(0.9, 0.98), eps=1e-9))

    train(data, model, criterion, optimizer)
    print(f"<<<<<<< finished train, cost {time.time() - train_start:.4f} seconds")

    # 预测
    # 加载模型
    model.load_state_dict(torch.load(SAVE_FILE))
    # 开始预测
    print(">>>>>>> start evaluate")
    evaluate_start = time.time()
    evaluate(data, model)
    print(f"<<<<<<< finished evaluate, cost {time.time() - evaluate_start:.4f} seconds")
